from pocketflow import Flow
from .nodes import (
    FreeConversationNode,
    CorrectionAnalysisNode,
    CorrectionReviewNode,
    ResponseFormattingNode
)

def create_chat_flow():
    """
    Create a flow for chat mode (free conversation).
    """
    # Create nodes
    free_conversation_node = FreeConversationNode()
    response_formatting_node = ResponseFormattingNode()
    
    # Connect nodes
    free_conversation_node >> response_formatting_node
    
    # Create flow
    return Flow(start=free_conversation_node)

def create_check_flow():
    """
    Create a flow for check mode (correction).
    """
    # Create nodes
    correction_analysis_node = CorrectionAnalysisNode()
    correction_review_node = CorrectionReviewNode()
    response_formatting_node = ResponseFormattingNode()
    
    # Connect nodes
    correction_analysis_node >> correction_review_node >> response_formatting_node
    
    # Create flow
    return Flow(start=correction_analysis_node)

# Create separate flows for each mode
chat_flow = create_chat_flow()
check_flow = create_check_flow()

# Function to get the appropriate flow based on mode
def get_flow(mode):
    """
    Return the appropriate flow based on the specified mode.
    
    Args:
        mode (str): The mode, either 'chat' or 'check'
        
    Returns:
        Flow: The appropriate flow for the specified mode
        
    Raises:
        ValueError: If the mode is not supported
    """
    if mode == "chat":
        return chat_flow
    elif mode == "check":
        return check_flow
    else:
        raise ValueError(f"Invalid mode: {mode}. Must be 'chat' or 'check'.")

# For backward compatibility, create a default flow
onceone_agent_flow = chat_flow